#include "linear/linear_hmm_engine.hpp"
#include "matrix_util.hpp"
#include "powerof2/linear_p2_crypto_agent.hpp"

#include <gtest/gtest.h>

bool only_multi_thread;
bool only_single_thread;

// Defualt BGV Params Set
params<helib::BGV> my_params;

bool should_run_test_case(const std::string &test_case_name) {
  if (only_single_thread) {
    return test_case_name == "_ST";
  } else if (only_multi_thread) {
    return test_case_name == "_MT";
  } else {
    return true;
  }
}

#define P2_HHMMSQ_TEST_TEMPLATE(DIM, THREAD)                                   \
  TEST(TEST_MNK_##DIM, _##THREAD) {                                            \
    /* Get the test case name  */                                              \
    const ::testing::TestInfo *const test_info =                               \
        ::testing::UnitTest::GetInstance()->current_test_info();               \
    const std::string test_case_name = test_info->name();                      \
                                                                               \
    if (!should_run_test_case(test_case_name)) {                               \
      GTEST_SKIP();                                                            \
    }                                                                          \
                                                                               \
    long m = DIM;                                                              \
    long n = DIM;                                                              \
    long k = DIM;                                                              \
                                                                               \
    NTL::mat_ZZ ptxt_A;                                                        \
    NTL::mat_ZZ ptxt_B;                                                        \
    FIX_NTL_RAND_SEED(true);                                                   \
    RAND_NTL_MATRIX(ptxt_A, m, k, 8);                                          \
    RAND_NTL_MATRIX(ptxt_B, k, n, 8);                                          \
    auto C = ptxt_A * ptxt_B;                                                  \
                                                                               \
    /* Create an inactive HMM engine */                                        \
    hmm_status<shmm_engine> hmm_status(NAIVE, test_case_name == "_ST");        \
    linear_hmme hmme(hmm_status);                                              \
                                                                               \
    /* Initialize HE cryptograph system */                                     \
    const auto eh_mnk = linear_hmme::get_expected_mnk(m, n, k);                \
    linearp2_hmmcc hmmcc(my_params, linear_hmme::get_expected_nslot(eh_mnk));  \
    hmmcc.report_context();                                                    \
                                                                               \
    /* Activate the HMM engine using EA */                                     \
    hmme.register_engine(hmmcc.context()->getEA(), eh_mnk);                    \
                                                                               \
    auto ctxt_A = hmmcc.encrypt(ptxt_A);                                       \
    auto ctxt_B = hmmcc.encrypt(ptxt_B);                                       \
                                                                               \
    auto ctxt_C = hmmcc.get_empty_ctxt();                                      \
    auto matrix_type = shmm_engine::parse_matrix_type(eh_mnk);                 \
    HELIB_NTIMER_START(_HMM);                                                  \
    hmme.cgemm(matrix_type, ctxt_C, ctxt_A, ctxt_B);                           \
    HELIB_NTIMER_STOP(_HMM);                                                   \
    NTL::mat_ZZ ptxt_C;                                                        \
    hmmcc.decrypt(ptxt_C, ctxt_C);                                             \
                                                                               \
    std::cout << "===================== HMM Report ======================\n";  \
    hmmcc.report_time();                                                       \
    std::cout << "* HMM Time (Outside)         : "                             \
              << _named_local_timer_HMM.getTime() << " s\n";                   \
    hmme.report();                                                             \
    std::cout << "=======================================================\n";  \
                                                                               \
    for (auto i = 0; i < m; i++) {                                             \
      for (auto j = 0; j < n; j++) {                                           \
        auto expected = C[i][j];                                               \
        auto actual = ptxt_C[0][i * n + j];                                    \
        EXPECT_EQ(expected, actual)                                            \
            << " diff at index (" << i << ", " << j << ")";                    \
      }                                                                        \
    }                                                                          \
  }

#define P2_HHMMSQ_TEST(DIM)                                                    \
  P2_HHMMSQ_TEST_TEMPLATE(DIM, ST)                                             \
  P2_HHMMSQ_TEST_TEMPLATE(DIM, MT)

P2_HHMMSQ_TEST(4)
P2_HHMMSQ_TEST(8)
P2_HHMMSQ_TEST(16)
P2_HHMMSQ_TEST(32)
P2_HHMMSQ_TEST(64)
P2_HHMMSQ_TEST(128)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);

  // Create a `command_line_parser` using `params<helib::BGV>`
  command_line_parser parser(my_params);
  parser.add_option("-mt", "only_multi_thread", &only_multi_thread);
  parser.add_option("-st", "only_single_thread", &only_single_thread);

  //* Parse command line
  int ret = parser.parse(argc, argv);
  EXPECT_EQ(ret, 0);

  ret = RUN_ALL_TESTS();

  return ret;
}